 % Simulation
 
clc
clear all
close all

% calibrated beta and d_I
beta=0.15246;
d_I=0.01634;

w=ones(26,2); % no wage differentials
load moments_from_Seoul_survey.mat
load distance_GIS.mat
load calibrated_E.mat
load initial_infection.mat

epsilon_weekday=4.1642;
epsilon_weekend=4.9144;
kappa=0.0339;

H_R(:,1)=H_R_young; % young
H_R(:,2)=H_R_old;   % old

ncity=25;

gamma=1/18; % 18 days to be recovered
tau=[1/8.5 1/10.2]; % days to be quaurantined
death_rate=[0.0021 0.0273]; % fatality rate

delta_t(:,1)=[0.00583956102 0.00583956102]'; % young, weekday weekend
delta_t(:,2)=[0.00736463186 0.00736463186]'; % old, weekday weekend

delta_j(:,1)=[0.00209103789 0.00138626867 0.00209103789 0.00138626867]'; % young, weekday case/visit weekend case/visit
delta_j(:,2)=[0.00247622592 0.00096650016 0.00247622592 0.00096650016]'; % old, weekday case/visit weekend case/visit
%delta_j(1,:)=zeros(1,2); delta_j(3,:)=zeros(1,2); % partial disclosure
%delta_j(2,:)=zeros(1,2); delta_j(4,:)=zeros(1,2); % partial disclosure
%delta_j=zeros(4,2); % no disclosure

global w d_GIS ncity epsilon_weekday epsilon_weekend kappa E_weekday E_weekend H_R initial_infection gamma tau death_rate beta d_I delta_j delta_t

time=1000; period=[1:1:time];

%% Actual Simulation
% setup dates: starting from Jan 30
Day = datetime(2020,1,30) + caldays(0:time);

% Initial Values
    Qtime(1,:,1)=initial_infection'; % Initial detected, as of January 30 (all looks exogenous)
    % city 8(1), 13(1), 15(1), 16(1) // 4 the detected
    Qtime(1,:,2)=zeros(1,25);
    Itime(1,:,1)=9*initial_infection';
    Itime(1,:,2)=zeros(1,25);
    for a=1:2
        for d1=1:ncity
            Ntime(1,d1,a)=H_R(d1,a);
            Stime(1,d1,a)=Ntime(1,d1,a)-Itime(1,d1,a)-Qtime(1,d1,a);
            Rtime(1,d1,a)=Ntime(1,d1,a)-Itime(1,d1,a)-Stime(1,d1,a)-Qtime(1,d1,a); % zero
            Ftime(1,d1,a)=death_rate(a)*Rtime(1,d1,a);
            Dtime(1,d1,a)=0;
            cum_Dtime(1,d1)=Qtime(1,d1,1)+Qtime(1,d1,2);
            cum_Dtime_14(1,d1)=Qtime(1,d1,1)+Qtime(1,d1,2);
        end
    end

% run simulation
for t=2:time
    [DayNumber,DayName] = weekday(Day(t));
    if DayNumber>=2 && DayNumber<=6 % weekdays
            [pi(t,:,:,:) d_ija(t,:,:,:)]=Calculate_pi_function(cum_Dtime_14(t-1,:));
                for d1=1:ncity
                    for a=1:2
                    S_dot(t,d1,a)=0;
                    for d2=1:ncity
                        S_dot(t,d1,a)=S_dot(t,d1,a)-beta*(sum(pi(t,:,d2,1).*Itime(t-1,:,1))+sum(pi(t,:,d2,2).*Itime(t-1,:,2)))/(sum(pi(t,:,d2,1).*(Ntime(t-1,:,1)-Qtime(t-1,:,1)))+sum(pi(t,:,d2,2).*(Ntime(t-1,:,2)-Qtime(t-1,:,2))))*(pi(t,d1,d2,a)*Stime(t-1,d1,a));
                    end
                    Stime(t,d1,a) =  Stime(t-1,d1,a) + S_dot(t,d1,a);
                    Itime(t,d1,a) =  Itime(t-1,d1,a)*(1 - gamma) - S_dot(t,d1,a) - d_I*Itime(t-1,d1,a);
                    Qtime(t,d1,a) =  Qtime(t-1,d1,a)*(1 - tau(a)) + d_I*Itime(t-1,d1,a);
                    Rtime(t,d1,a) =  Rtime(t-1,d1,a) + Itime(t-1,d1,a)*gamma + tau(a)*Qtime(t-1,d1,a);
                    Ftime(t,d1,a) =  death_rate(a)*Rtime(t,d1,a);
                    Ntime(t,d1,a) =  H_R(d1,a) - death_rate(a)*Rtime(t,d1,a);
                    Dtime(t,d1,a) = d_I*Itime(t-1,d1,a);
                    end
                    cum_Dtime(t,d1) = cum_Dtime(t-1,d1) + d_I*Itime(t-1,d1,1) + d_I*Itime(t-1,d1,2); % adding newly detected cases
                    cum_Dtime_14(t,d1) = sum(Dtime(max(1,t-13):t,d1,1))+sum(Dtime(max(1,t-13):t,d1,2));
                end
    else % weekends
            [pi(t,:,:,:) d_ija(t,:,:,:)]=Calculate_pi_weekend_function(cum_Dtime_14(t-1,:));
                for d1=1:ncity
                    for a=1:2
                    S_dot(t,d1,a)=0;
                    for d2=1:ncity
                        S_dot(t,d1,a)=S_dot(t,d1,a)-beta*(sum(pi(t,:,d2,1).*Itime(t-1,:,1))+sum(pi(t,:,d2,2).*Itime(t-1,:,2)))/(sum(pi(t,:,d2,1).*(Ntime(t-1,:,1)-Qtime(t-1,:,1)))+sum(pi(t,:,d2,2).*(Ntime(t-1,:,2)-Qtime(t-1,:,2))))*(pi(t,d1,d2,a)*Stime(t-1,d1,a));
                    end
                    Stime(t,d1,a) =  Stime(t-1,d1,a) + S_dot(t,d1,a);
                    Itime(t,d1,a) =  Itime(t-1,d1,a)*(1 - gamma) - S_dot(t,d1,a) - d_I*Itime(t-1,d1,a);
                    Qtime(t,d1,a) =  Qtime(t-1,d1,a)*(1 - tau(a)) + d_I*Itime(t-1,d1,a);
                    Rtime(t,d1,a) =  Rtime(t-1,d1,a) + Itime(t-1,d1,a)*gamma + tau(a)*Qtime(t-1,d1,a);
                    Ftime(t,d1,a) =  death_rate(a)*Rtime(t,d1,a);
                    Ntime(t,d1,a) =  H_R(d1,a) - death_rate(a)*Rtime(t,d1,a);
                    Dtime(t,d1,a) = d_I*Itime(t-1,d1,a);
                    end
                    cum_Dtime(t,d1) = cum_Dtime(t-1,d1) + d_I*Itime(t-1,d1,1) + d_I*Itime(t-1,d1,2); % adding newly detected cases
                    cum_Dtime_14(t,d1) = sum(Dtime(max(1,t-13):t,d1,1))+sum(Dtime(max(1,t-13):t,d1,2));
                end
    end
end

cum_Dtime_in_two_month=cum_Dtime(123,:)';
target1=sum(cum_Dtime_in_two_month) % 861

Itime_sum=Itime(:,:,1)+Itime(:,:,2);
Qtime_sum=Qtime(:,:,1)+Qtime(:,:,2);
fraction=sum(Itime_sum,2)./(sum(Qtime_sum,2)+sum(Itime_sum,2)); % Our primary estimates for the fraction of infections that are undetected range from 88.7% to 93.6%. (Iceland)
target2=mean(fraction(1:123)) % new target 90%

day=period;
Pop=sum(H_R);
Stime_sum=Stime(:,:,1)+Stime(:,:,2);
Rtime_sum=Rtime(:,:,1)+Rtime(:,:,2);
Ftime_sum=Ftime(:,:,1)+Ftime(:,:,2);

%% Figure 3a
cum_Dtime_in_two_month=cum_Dtime(123,:)'; % as of May 31st
load infection_May_31st.mat
figure(1)
x = linspace(0,80);
y = linspace(0,80);
plot(x,y,'k');
hold on
scatter(cum_Dtime_in_two_month,post_infection,'k');
xh=xlabel('model','FontSize',14);
yh=ylabel('data','FontSize',14);
hold off
saveas(gcf,'Output/Figure3a.png')

%% Figure 3b
figure(2)
plot(day,sum(Itime(:,:,1),2)/Pop(1), 'k' , 'LineWidth', 2); hold on;
plot(day,sum(Itime(:,:,2),2)/Pop(2), '--k' , 'LineWidth', 2); hold on;
lh=legend({'Age 20-59', 'Age 60+'},'Location','northeast');set(lh,'fontsize',12);legend boxoff
xh=xlabel('Day','FontSize',14); 
yh=ylabel('Share of Population','FontSize',14); 
saveas(gcf,'Output/Figure3b.png')

%% Figure 3c
% calculate model-predicted total inflow by district
model_flow_weekdays_sum=zeros(time,25,2);
model_flow_weekends_sum=zeros(time,25,2);
for t=2:time
    [DayNumber,DayName] = weekday(Day(t));
    if DayNumber>=2 && DayNumber<=6 % weekdays
        for a=1:2
            for d=1:25
                for o=1:25
                        model_flow_weekdays_sum(t,d,a)=model_flow_weekdays_sum(t,d,a)+pi(t,o,d,a)*(H_R(o,a)-Qtime(t,o,a)-death_rate(a)*Rtime(t,o,a));
                end
            end
        end
    else % weekends
        for a=1:2
            for d=1:25
                for o=1:25
                        model_flow_weekends_sum(t,d,a)=model_flow_weekends_sum(t,d,a)+pi(t,o,d,a)*(H_R(o,a)-Qtime(t,o,a)-death_rate(a)*Rtime(t,o,a));
                end
            end
        end
    end
end
size_weekdays=size(model_flow_weekdays_sum);
for a=1:2
for t=1:size_weekdays(1)
    for d=1:25
    if model_flow_weekdays_sum(t,d,a)==0
        model_flow_weekdays_sum(t,d,a)=NaN;
    end
    end
end
end
size_weekends=size(model_flow_weekends_sum);
for a=1:2
for t=1:size_weekends(1)
    for d=1:25
    if model_flow_weekends_sum(t,d,a)==0
        model_flow_weekends_sum(t,d,a)=NaN;
    end
    end
end
end
model_flow_weekdays_norm_avg=zeros(size_weekdays(1),25);
model_flow_weekends_norm_avg=zeros(size_weekends(1),25);
for d=1:25
    model_flow_weekdays_norm(:,d,1)=log(model_flow_weekdays_sum(:,d,1)./model_flow_weekdays_sum(2,d,1));
    model_flow_weekdays_norm(:,d,2)=log(model_flow_weekdays_sum(:,d,2)./model_flow_weekdays_sum(2,d,2));
    model_flow_weekdays_norm_avg(:,d)=model_flow_weekdays_norm(:,d,1)*(sum(H_R(d,1))/(sum(H_R(d,1))+sum(H_R(d,2))))+model_flow_weekdays_norm(:,d,2)*(sum(H_R(d,2))/(sum(H_R(d,1))+sum(H_R(d,2))));
    model_flow_weekends_norm(:,d,1)=log(model_flow_weekends_sum(:,d,1)./model_flow_weekends_sum(3,d,1));
    model_flow_weekends_norm(:,d,2)=log(model_flow_weekends_sum(:,d,2)./model_flow_weekends_sum(3,d,2));
    model_flow_weekends_norm_avg(:,d)=model_flow_weekends_norm(:,d,1)*(sum(H_R(d,1))/(sum(H_R(d,1))+sum(H_R(d,2))))+model_flow_weekends_norm(:,d,2)*(sum(H_R(d,2))/(sum(H_R(d,1))+sum(H_R(d,2))));
end
model_flow_weekdays_norm_avg_avg=mean(model_flow_weekdays_norm_avg,2);
model_flow_weekends_norm_avg_avg=mean(model_flow_weekends_norm_avg,2);
model_flow_weekdays_norm_avg_p95t=model_flow_weekdays_norm_avg_avg+1.96*std(model_flow_weekdays_norm_avg,0,2);
model_flow_weekdays_norm_avg_p95b=model_flow_weekdays_norm_avg_avg-1.96*std(model_flow_weekdays_norm_avg,0,2);
model_flow_weekends_norm_avg_p95t=model_flow_weekends_norm_avg_avg+1.96*std(model_flow_weekends_norm_avg,0,2);
model_flow_weekends_norm_avg_p95b=model_flow_weekends_norm_avg_avg-1.96*std(model_flow_weekends_norm_avg,0,2);
figure(3)
plot(day(1:size_weekdays),model_flow_weekdays_norm_avg_avg, 'k' , 'LineWidth', 6); hold on;
plot(day(1:size_weekends),model_flow_weekends_norm_avg_avg, 'k' , 'LineWidth', 4); hold on;
plot(day(1:size_weekdays),model_flow_weekdays_norm_avg_p95t, 'k' , 'LineWidth', 0.6); hold on;
plot(day(1:size_weekdays),model_flow_weekdays_norm_avg_p95b, 'k' , 'LineWidth', 0.6); hold on;
plot(day(1:size_weekends),model_flow_weekends_norm_avg_p95t, 'k' , 'LineWidth', 0.4); hold on;
plot(day(1:size_weekends),model_flow_weekends_norm_avg_p95b, 'k' , 'LineWidth', 0.4); hold on;
lh=legend({'weekdays','weekends'},'Location','southeast');set(lh,'fontsize',12);legend boxoff
xh=xlabel('Day','FontSize',14); 
yh=ylabel('Inflows by District (log scale)','FontSize',14);
saveas(gcf,'Output/Figure3c.png')

%% Figure 3d
common_z_weekday=draw_frechet(E_weekday,epsilon_weekday,10^5);
common_z_weekend=draw_frechet(E_weekend,epsilon_weekend,10^5);

global common_z_weekday common_z_weekend
for t=1:time
    [DayNumber,DayName] = weekday(Day(t));
    if DayNumber>=2 && DayNumber<=6 % weekdays
        utility_weekdays(t,:,:)=Calculate_utility_per_capita_function(cum_Dtime_14(t,:),squeeze(Qtime(t,:,:)),squeeze(Rtime(t,:,:)));
    else
        utility_weekends(t,:,:)=Calculate_utility_per_capita_weekend_function(cum_Dtime_14(t,:),squeeze(Qtime(t,:,:)),squeeze(Rtime(t,:,:)));
    end
    t
end

utility_weekdays_sum(:,1)=sum(utility_weekdays(:,:,1),2); % per capita
utility_weekdays_sum(:,2)=sum(utility_weekdays(:,:,2),2); % per capita
size_weekdays=size(utility_weekdays_sum);
for a=1:2
for t=1:size_weekdays(1)
    if utility_weekdays_sum(t,a)==0
        utility_weekdays_sum(t,a)=NaN;
    end
end
end
utility_weekdays_sum_norm(:,1)=log(utility_weekdays_sum(:,1)./utility_weekdays_sum(1,1));
utility_weekdays_sum_norm(:,2)=log(utility_weekdays_sum(:,2)./utility_weekdays_sum(1,2));

utility_weekends_sum(:,1)=sum(utility_weekends(:,:,1),2); % per capita
utility_weekends_sum(:,2)=sum(utility_weekends(:,:,2),2); % per capita
size_weekends=size(utility_weekends_sum);
for a=1:2
for t=1:size_weekends(1)
    if utility_weekends_sum(t,a)==0
        utility_weekends_sum(t,a)=NaN;
    end
end
end
utility_weekends_sum_norm(:,1)=log(utility_weekends_sum(:,1)./utility_weekends_sum(3,1)); % first weekend, day 3
utility_weekends_sum_norm(:,2)=log(utility_weekends_sum(:,2)./utility_weekends_sum(3,2)); % first weekend, day 3

utility_weekdays_sum_norm_avg=utility_weekdays_sum_norm(:,1)*(Pop(1)/(Pop(1)+Pop(2)))+utility_weekdays_sum_norm(:,2)*(Pop(2)/(Pop(1)+Pop(2)));
utility_weekends_sum_norm_avg=utility_weekends_sum_norm(:,1)*(Pop(1)/(Pop(1)+Pop(2)))+utility_weekends_sum_norm(:,2)*(Pop(2)/(Pop(1)+Pop(2)));

for t=2:time
    [DayNumber,DayName] = weekday(Day(t));
    if DayNumber>=2 && DayNumber<=6 % weekdays
        utility_sum_norm_avg(t,1)=utility_weekdays_sum_norm_avg(t);
        utility_sum_norm_young(t,1)=utility_weekdays_sum_norm(t,1);
        utility_sum_norm_old(t,1)=utility_weekdays_sum_norm(t,2);
    else
        utility_sum_norm_avg(t,1)=utility_weekends_sum_norm_avg(t);
        utility_sum_norm_young(t,1)=utility_weekends_sum_norm(t,1);
        utility_sum_norm_old(t,1)=utility_weekends_sum_norm(t,2);
    end
end
utility_sum_norm_MA = movmean(utility_sum_norm_avg,7);
figure(4)
plot(day(2:time),utility_sum_norm_MA(2:time), 'k' , 'LineWidth', 3); hold on;
xh=xlabel('Day','FontSize',14); 
yh=ylabel('Economic Welfare (log scale, 7-days m.a.)','FontSize',14);
saveas(gcf,'Output/Figure3d.png')

%% Save files
total_cum_Dtime=sum(cum_Dtime,2);
save ('Output/cnt_original.mat', 'Day', 'day', 'Pop', 'Stime', 'Itime', 'Qtime', 'Rtime', 'Ftime', 'total_cum_Dtime', 'utility_sum_norm_avg', 'utility_sum_norm_young','utility_sum_norm_old')
%save ('Output/cnt_delta1_0.mat', 'day', 'Pop', 'Stime', 'Itime', 'Qtime', 'Rtime', 'Ftime', 'total_cum_Dtime', 'utility_sum_norm_avg', 'utility_sum_norm_young','utility_sum_norm_old')
%save ('Output/cnt_delta2_0.mat', 'day', 'Pop', 'Stime', 'Itime', 'Qtime', 'Rtime', 'Ftime', 'total_cum_Dtime', 'utility_sum_norm_avg', 'utility_sum_norm_young','utility_sum_norm_old')
%save ('Output/cnt_delta1_0_delta2_0.mat', 'day', 'Pop', 'Stime', 'Itime', 'Qtime', 'Rtime', 'Ftime', 'total_cum_Dtime', 'utility_sum_norm_avg', 'utility_sum_norm_young','utility_sum_norm_old')
